#include <opencv2/opencv.hpp>
#include <cmath>
#include <iostream>

using namespace cv;
using namespace std;

// XYZ颜色空间转换矩阵
const double Mrgb[3][3] = {
    {0.4124564, 0.3575761, 0.1804375},
    {0.2126729, 0.7151522, 0.0721750},
    {0.0193339, 0.1191920, 0.9503041}
};

const double Mxyz[3][3] = {
    { 3.2404542, -1.5371385, -0.4985314},
    {-0.9692660,  1.8760108,  0.0415560},
    { 0.0556434, -0.2040259,  1.0572252}
};

double Transform(double x) {
    return x > 0.0031308 ? 1.055*pow(x,1/2.4)-0.055 : 12.92*x;
}

Mat Tone_Mapping(const string& fileName, double bias=0.85, double C=1.0, double r=0.6) {
    Mat hdrImg;
    bool isHDR = false;
    
    // 读取HDR图像（支持.hdr和.exr格式）
    if (fileName.find(".hdr") != string::npos || fileName.find(".exr") != string::npos) {
        hdrImg = imread(fileName, IMREAD_ANYDEPTH | IMREAD_COLOR);
        isHDR = true;
        if(hdrImg.empty()) {
            cerr << "Error loading HDR image!" << endl;
            return Mat();
        }
    } else {
        hdrImg = imread(fileName, IMREAD_COLOR);
        if(hdrImg.empty()) {
            cerr << "Error loading image!" << endl;
            return Mat();
        }
        hdrImg.convertTo(hdrImg, CV_32FC3, 1/255.0);
    }

    const int rows = hdrImg.rows;
    const int cols = hdrImg.cols;
    Mat out_img(rows, cols, CV_8UC3);

    double Lsum = 0;
    double Lmax = 0;
    const double a = 0.0001;

    // 第一遍遍历：计算亮度统计量
    for(int i=0; i<rows; ++i) {
        for(int j=0; j<cols; ++j) {
            Vec3f pixel = hdrImg.at<Vec3f>(i,j);
            double R = pixel[2], G = pixel[1], B = pixel[0];
            
            // 转换到XYZ颜色空间
            double X = Mrgb[0][0]*R + Mrgb[0][1]*G + Mrgb[0][2]*B;
            double Y = Mrgb[1][0]*R + Mrgb[1][1]*G + Mrgb[1][2]*B;
            double Z = Mrgb[2][0]*R + Mrgb[2][1]*G + Mrgb[2][2]*B;

            // 计算亮度统计
            double Lw = Y;
            Lsum += log(Lw + a);
            Lmax = max(Lmax, Lw);
        }
    }

    double Lav = exp(Lsum / (rows*cols));
    Lmax /= Lav;

    // 第二遍遍历：应用色调映射
    for(int i=0; i<rows; ++i) {
        for(int j=0; j<cols; ++j) {
            Vec3f pixel = hdrImg.at<Vec3f>(i,j);
            double R = pixel[2], G = pixel[1], B = pixel[0];
            
            // RGB转XYZ
            double X = Mrgb[0][0]*R + Mrgb[0][1]*G + Mrgb[0][2]*B;
            double Y = Mrgb[1][0]*R + Mrgb[1][1]*G + Mrgb[1][2]*B;
            double Z = Mrgb[2][0]*R + Mrgb[2][1]*G + Mrgb[2][2]*B;

            // 归一化处理
            double sum = X + Y + Z;
            double x = X/sum, y = Y/sum;

            // 核心色调映射公式
            double L = Y / Lav;
            double Ld = log(L + 1) / log(2 + 8*pow((L/Lmax), log(bias)/log(0.5))) / log10(Lmax + 1);

            // 恢复色度
            X = Ld * x / y;
            Z = Ld * (1 - x - y) / y;
            Y = Ld;

            // XYZ转RGB
            double R_new = Mxyz[0][0]*X + Mxyz[0][1]*Y + Mxyz[0][2]*Z;
            double G_new = Mxyz[1][0]*X + Mxyz[1][1]*Y + Mxyz[1][2]*Z;
            double B_new = Mxyz[2][0]*X + Mxyz[2][1]*Y + Mxyz[2][2]*Z;

            // 伽马校正
            R_new = C * pow(max(R_new,0.0), r);
            G_new = C * pow(max(G_new,0.0), r);
            B_new = C * pow(max(B_new,0.0), r);

            // 颜色修正
            R_new = Transform(R_new);
            G_new = Transform(G_new);
            B_new = Transform(B_new);

            // 写入输出图像
            out_img.at<Vec3b>(i,j) = Vec3b(
                saturate_cast<uchar>(B_new*255),
                saturate_cast<uchar>(G_new*255),
                saturate_cast<uchar>(R_new*255)
            );
        }
    }

    return out_img;
}

int main() {
    Mat result = Tone_Mapping("input.hdr", 0.85, 1.0, 0.6);
    imwrite("output.jpg", result);
    imshow("Result", result);
    waitKey(0);
    return 0;
}